# -*- coding: utf-8 -*-
"""
沪深300ETF网格交易策略回测系统
依赖库：pip install akshare pandas matplotlib
"""

import akshare as ak
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

# 设置中文字体（任选其一）
plt.rcParams['font.sans-serif'] = ['SimHei']  # Windows系统黑体
# plt.rcParams['font.sans-serif'] = ['Arial Unicode MS']  # macOS
# plt.rcParams['font.sans-serif'] = ['WenQuanYi Zen Hei']  # Linux

# 解决负号显示问题
plt.rcParams['axes.unicode_minus'] = False

# ====================
# 策略参数配置
# ====================
STRATEGY_PARAMS = {
    'initial_cash': 20000,  # 初始资金
    'reserve_ratio': 0.2,  # 现金预留比例
    'grid_spacing': 0.005,  # 网格间距(0.5%)
    'shares_per_grid': 500,  # 每格交易股数
    'price_range': [3.70, 3.95],  # 价格区间
    'commission': 0.0001,  # 手续费率(万1)
    'benchmark_price': 3.85,  # 基准价
    'symbol': "510310",  # ETF代码
    'start_date': "2024-10-15",  # 回测开始日期
    'end_date': "2025-03-28"  # 回测结束日期
}


# ====================
# 数据获取模块
# ====================
def fetch_historical_data(params):
    """从AKShare获取ETF历史数据"""
    print("\n正在获取历史数据...")
    try:
        # 获取基金历史数据
        df = ak.fund_etf_hist_em(symbol=params['symbol'], period="daily")
        df['日期'] = pd.to_datetime(df['日期'])

        # 过滤日期范围
        mask = (df['日期'] >= params['start_date']) & (df['日期'] <= params['end_date'])
        df = df[mask].set_index('日期').sort_index()

        # 数据校验
        if df.empty:
            raise ValueError("获取到的数据为空，请检查日期范围")

        print(f"成功获取{len(df)}个交易日数据")
        print(f"最新数据日期：{df.index[-1].strftime('%Y-%m-%d')}")
        return df['收盘']

    except Exception as e:
        print(f"数据获取失败: {str(e)}")
        print("可能原因：\n1. AKShare需要升级(pip install akshare --upgrade)"
              "\n2. 基金代码错误\n3. 日期范围超出数据范围")
        return None


# ====================
# 网格交易引擎
# ====================
class GridTradingSystem:
    def __init__(self, params):
        self.params = params
        self.reset()

    def reset(self):
        """重置回测状态"""
        self.cash = self.params['initial_cash'] * (1 - self.params['reserve_ratio'])
        self.reserve_cash = self.params['initial_cash'] * self.params['reserve_ratio']
        self.hold_shares = 0
        self.last_trade_price = self.params['benchmark_price']
        self.trade_log = []
        self.position_history = []

    def execute_order(self, price, direction):
        """执行交易订单"""
        amount = price * self.params['shares_per_grid']
        commission = amount * self.params['commission']

        if direction == 'buy':
            if self.cash >= (amount + commission):
                self.cash -= (amount + commission)
                self.hold_shares += self.params['shares_per_grid']
                return True
        elif direction == 'sell':
            if self.hold_shares >= self.params['shares_per_grid']:
                self.cash += (amount - commission)
                self.hold_shares -= self.params['shares_per_grid']
                return True
        return False

    def run_backtest(self, price_series):
        """运行回测引擎"""
        print("\n开始回测...")
        for date, price in price_series.items():
            # 跳过无效价格
            if np.isnan(price):
                continue

            # 计算价格波动
            price_change = (price - self.last_trade_price) / self.last_trade_price

            # 触发交易逻辑
            if price_change >= self.params['grid_spacing']:
                if self.execute_order(price, 'sell'):
                    self.trade_log.append({'date': date, 'price': price, 'action': 'sell'})
                    self.last_trade_price = price
            elif price_change <= -self.params['grid_spacing']:
                if self.execute_order(price, 'buy'):
                    self.trade_log.append({'date': date, 'price': price, 'action': 'buy'})
                    self.last_trade_price = price

            # 记录持仓
            self.position_history.append({
                'date': date,
                'total_assets': self.cash + self.hold_shares * price + self.reserve_cash
            })

        print("回测完成！")
        return pd.DataFrame(self.trade_log), pd.DataFrame(self.position_history)


# ====================
# 结果分析模块
# ====================
def analyze_results(trade_df, position_df, params, prices):
    """分析回测结果"""
    # 计算关键指标
    final_assets = position_df.iloc[-1]['total_assets']
    total_profit = final_assets - params['initial_cash']
    return_rate = total_profit / params['initial_cash']

    # 打印报告
    print("\n" + "=" * 40)
    print("回测结果摘要")
    print("=" * 40)
    print(f"期初资金: {params['initial_cash']:,.2f}元")
    print(f"期末总资产: {final_assets:,.2f}元")
    print(f"绝对收益: {total_profit:+,.2f}元")
    print(f"收益率: {return_rate * 100:+.2f}%")
    print(f"交易次数: {len(trade_df)}次")
    print(f"买入次数: {len(trade_df[trade_df['action'] == 'buy'])}次")
    print(f"卖出次数: {len(trade_df[trade_df['action'] == 'sell'])}次")

    # 可视化结果
    plt.figure(figsize=(12, 8))

    # 价格和交易点
    ax1 = plt.subplot(2, 1, 1)
    prices.plot(ax=ax1, label='ETF价格', color='navy')
    buy_dates = trade_df[trade_df['action'] == 'buy']['date']
    ax1.scatter(buy_dates, trade_df[trade_df['action'] == 'buy']['price'],
                marker='^', color='green', label='买入点')
    sell_dates = trade_df[trade_df['action'] == 'sell']['date']
    ax1.scatter(sell_dates, trade_df[trade_df['action'] == 'sell']['price'],
                marker='v', color='red', label='卖出点')
    ax1.set_title('价格走势与交易信号')
    ax1.legend()

    # 资产曲线
    ax2 = plt.subplot(2, 1, 2)
    position_df.set_index('date')['total_assets'].plot(ax=ax2, color='darkorange')
    ax2.set_title('账户总资产变化')
    ax2.set_ylabel('资产总额（元）')

    plt.tight_layout()
    plt.show()

    # 保存结果
    trade_df.to_csv('../data/output/grid_trading_log.csv', index=False)
    position_df.to_csv('../data/output/asset_history.csv', index=False)
    print("\n交易记录已保存至 grid_trading_log.csv")


# ====================
# 主程序
# ====================
if __name__ == "__main__":
    # 获取数据
    price_data = fetch_historical_data(STRATEGY_PARAMS)
    if price_data is None:
        exit()

    # 初始化交易系统
    trading_system = GridTradingSystem(STRATEGY_PARAMS)

    # 运行回测
    trades, positions = trading_system.run_backtest(price_data)

    # 分析结果
    analyze_results(trades, positions, STRATEGY_PARAMS, price_data)